package uestc.lj.basicPromotion.union;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * 并查集
 *
 * @Author:Crazlee
 * @Date:2021/11/26
 */
public class Code01_UnionFind {

	public static class Element<V> {
		public V value;

		public Element(V value) {
			this.value = value;
		}
	}

	/**
	 * 1
	 * / \
	 * 2   3
	 * \
	 * 4
	 *
	 * @param <V>
	 */
	public static class UnionFindSet<V> {
		/**
		 * 样本对应自己定义的元素表
		 * 1——①
		 * 2——②
		 * 3——③
		 * 4——④
		 */
		public HashMap<V, Element<V>> elementMap;

		/**
		 * key——元素   value——父元素
		 * ①——①
		 * ②——①
		 * ③——①
		 * ④——②
		 */
		public HashMap<Element<V>, Element<V>> fatherMap;

		/**
		 * key——某个集合的代表元素   value——该集合的大小
		 * ①——4
		 */
		public HashMap<Element<V>, Integer> sizeMap;

		/**
		 * 初始化
		 *
		 * @param list 用户给的初始值
		 */
		public UnionFindSet(List<V> list) {
			elementMap = new HashMap<>();
			fatherMap = new HashMap<>();
			sizeMap = new HashMap<>();
			for (V value : list) {
				Element<V> element = new Element<>(value);
				elementMap.put(value, element);
				fatherMap.put(element, element);
				sizeMap.put(element, 1);
			}
		}

		/**
		 * 给定一个element，向上寻找其代表元素
		 *
		 * @param element 待查元素
		 * @return 待查元素的父元素
		 */
		private Element<V> findHead(Element<V> element) {
			Stack<Element<V>> path = new Stack<>();
			// 1.寻找element的最顶部元素
			while (element != fatherMap.get(element)) {
				path.push(element);
				element = fatherMap.get(element);
			}
			// 2.扁平化处理
			while (!path.isEmpty()) {
				fatherMap.put(path.pop(), element);
			}
			return element;
		}


		/**
		 * 判断两个元素是否是同一集合
		 *
		 * @param a 元素a
		 * @param b 元素b
		 * @return 是否是同一集合
		 */
		public boolean isSameSet(V a, V b) {
			// 1.首先判断a和b是否在初始化时被注册到hash表中
			if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
				// 2.判断 a和b 的代表元素是否是同一个元素
				return findHead(elementMap.get(a)) == findHead(elementMap.get(b));
			}
			return false;
		}

		/**
		 * 将两个元素集合合并成一个元素集合
		 * 数量少的合并到数量大的
		 *
		 * @param a 元素a
		 * @param b 元素b
		 */
		public void union(V a, V b) {
			// 1.判断两个元素是否被注册进hashmap中
			if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
				// 2.判断两个元素集的代表元素是否相等
				Element<V> aHead = findHead(elementMap.get(a));
				Element<V> bHead = findHead(elementMap.get(b));
				if (aHead != bHead) {
					// 3.将数量小的挂在数量大的后面
					Element<V> big = sizeMap.get(aHead) > sizeMap.get(bHead) ? aHead : bHead;
					Element<V> small = big == aHead ? bHead : aHead;
					fatherMap.put(small, big);
					sizeMap.put(big, sizeMap.get(aHead) + sizeMap.get(bHead));
					sizeMap.remove(small);
				}
			}
		}
	}
}
